#!/usr/bin/env python3
# -*- coding: utf-8 -*-
import logging
import socket
import ssl
import time
import threading
import types
import errno

try:
  from OpenSSL.SSL import Error as pyOpenSSLError
except ImportError:

  class pyOpenSSLError(Exception):
    pass


from ambari_ws4py import WS_KEY, WS_VERSION
from ambari_ws4py.exc import HandshakeError, StreamClosed
from ambari_ws4py.streaming import Stream
from ambari_ws4py.messaging import Message, PingControlMessage, PongControlMessage
from ambari_ws4py.compat import basestring, unicode

DEFAULT_READING_SIZE = 2

logger = logging.getLogger("ambari_ws4py")

__all__ = ["WebSocket", "EchoWebSocket", "Heartbeat"]


class Heartbeat(threading.Thread):
  def __init__(self, websocket, frequency=2.0):
    """
    Runs at a periodic interval specified by
    `frequency` by sending an unsolicitated pong
    message to the connected peer.

    If the message fails to be sent and a socket
    error is raised, we close the websocket
    socket automatically, triggering the `closed`
    handler.
    """
    threading.Thread.__init__(self)
    self.websocket = websocket
    self.frequency = frequency

  def __enter__(self):
    if self.frequency:
      self.start()
    return self

  def __exit__(self, exc_type, exc_value, exc_tb):
    self.stop()

  def stop(self):
    self.running = False

  def run(self):
    self.running = True
    while self.running:
      time.sleep(self.frequency)
      if self.websocket.terminated:
        break

      try:
        self.websocket.send(PongControlMessage(data="beep"))
      except socket.error:
        logger.info("Heartbeat failed")
        self.websocket.server_terminated = True
        self.websocket.close_connection()
        break


class WebSocket(object):
  """Represents a websocket endpoint and provides a high level interface to drive the endpoint."""

  def __init__(
    self, sock, protocols=None, extensions=None, environ=None, heartbeat_freq=None
  ):
    """The ``sock`` is an opened connection
    resulting from the websocket handshake.

    If ``protocols`` is provided, it is a list of protocols
    negotiated during the handshake as is ``extensions``.

    If ``environ`` is provided, it is a copy of the WSGI environ
    dictionnary from the underlying WSGI server.
    """

    self.stream = Stream(always_mask=False)
    """
        Underlying websocket stream that performs the websocket
        parsing to high level objects. By default this stream
        never masks its messages. Clients using this class should
        set the ``stream.always_mask`` fields to ``True``
        and ``stream.expect_masking`` fields to ``False``.
        """

    self.protocols = protocols
    """
        List of protocols supported by this endpoint.
        Unused for now.
        """

    self.extensions = extensions
    """
        List of extensions supported by this endpoint.
        Unused for now.
        """

    self.sock = sock
    """
        Underlying connection.
        """

    self._is_secure = hasattr(sock, "_ssl") or hasattr(sock, "_sslobj")
    """
        Tell us if the socket is secure or not.
        """

    self.client_terminated = False
    """
        Indicates if the client has been marked as terminated.
        """

    self.server_terminated = False
    """
        Indicates if the server has been marked as terminated.
        """

    self.reading_buffer_size = DEFAULT_READING_SIZE
    """
        Current connection reading buffer size.
        """

    self.environ = environ
    """
        WSGI environ dictionary.
        """

    self.heartbeat_freq = heartbeat_freq
    """
        At which interval the heartbeat will be running.
        Set this to `0` or `None` to disable it entirely.
        """
    "Internal buffer to get around SSL problems"
    self.buf = b""

    self._local_address = None
    self._peer_address = None

    self.lock = threading.Lock()

    "Used to signal that the server side should be shutting down"
    self.server_terminate_request = False

  @property
  def local_address(self):
    """
    Local endpoint address as a tuple
    """
    if not self._local_address:
      self._local_address = self.sock.getsockname()
      if len(self._local_address) == 4:
        self._local_address = self._local_address[:2]
    return self._local_address

  @property
  def peer_address(self):
    """
    Peer endpoint address as a tuple
    """
    if not self._peer_address:
      self._peer_address = self.sock.getpeername()
      if len(self._peer_address) == 4:
        self._peer_address = self._peer_address[:2]
    return self._peer_address

  def opened(self):
    """
    Called by the server when the upgrade handshake
    has succeeded.
    """
    pass

  def close(self, code=1000, reason=""):
    """
    Call this method to initiate the websocket connection
    closing by sending a close frame to the connected peer.
    The ``code`` is the status code representing the
    termination's reason.

    Once this method is called, the ``server_terminated``
    attribute is set. Calling this method several times is
    safe as the closing frame will be sent only the first
    time.

    .. seealso:: Defined Status Codes http://tools.ietf.org/html/rfc6455#section-7.4.1
    """

    # If we are sending a fragmented frame with a generator this will make that stop
    self.server_terminate_request = True

    with self.lock:
      if not self.server_terminated:
        self.server_terminated = True
        try:
          self._write(
            self.stream.close(code=code, reason=reason).single(
              mask=self.stream.always_mask
            )
          )
        except Exception as ex:
          logger.error("Error when terminating the connection: %s", str(ex))

  def closed(self, code, reason=None):
    """
    Called  when the websocket stream and connection are finally closed.
    The provided ``code`` is status set by the other point and
    ``reason`` is a human readable message.

    .. seealso:: Defined Status Codes http://tools.ietf.org/html/rfc6455#section-7.4.1
    """
    pass

  @property
  def terminated(self):
    """
    Returns ``True`` if both the client and server have been
    marked as terminated.
    """
    return self.client_terminated is True and self.server_terminated is True

  @property
  def connection(self):
    return self.sock

  def close_connection(self):
    """
    Shutdowns then closes the underlying connection.
    """
    if self.sock:
      try:
        self.sock.shutdown(socket.SHUT_RDWR)
        self.sock.close()
      except:
        pass
      finally:
        self.sock = None

  def ping(self, message):
    """
    Send a ping message to the remote peer.
    The given `message` must be a unicode string.
    """
    self.send(PingControlMessage(message))

  def ponged(self, pong):
    """
    Pong message, as a :class:`messaging.PongControlMessage` instance,
    received on the stream.
    """
    pass

  def received_message(self, message):
    """
    Called whenever a complete ``message``, binary or text,
    is received and ready for application's processing.

    The passed message is an instance of :class:`messaging.TextMessage`
    or :class:`messaging.BinaryMessage`.

    .. note:: You should override this method in your subclass.
    """
    pass

  def unhandled_error(self, error):
    """
    Called whenever a socket, or an OS, error is trapped
    by ws4py but not managed by it. The given error is
    an instance of `socket.error` or `OSError`.

    Note however that application exceptions will not go
    through this handler. Instead, do make sure you
    protect your code appropriately in `received_message`
    or `send`.

    The default behaviour of this handler is to log
    the error with a message.
    """
    logger.exception("Failed to receive data")

  def _write(self, b):
    """
    Trying to prevent a write operation
    on an already closed websocket stream.

    This cannot be bullet proof but hopefully
    will catch almost all use cases.
    """
    if self.terminated or self.sock is None:
      raise RuntimeError("Cannot send on a terminated websocket")

    self.sock.sendall(b)

  def send(self, payload, binary=False):
    """
    Sends the given ``payload`` out.

    If ``payload`` is some bytes or a bytearray,
    then it is sent as a single message not fragmented.

    If ``payload`` is a generator, each chunk is sent as part of
    fragmented message.

    If ``binary`` is set, handles the payload as a binary message.
    """
    message_sender = self.stream.binary_message if binary else self.stream.text_message

    if isinstance(payload, basestring) or isinstance(payload, bytearray):
      m = message_sender(payload).single(mask=self.stream.always_mask)
      with self.lock:
        self._write(m)

    elif isinstance(payload, Message):
      data = payload.single(mask=self.stream.always_mask)
      with self.lock:
        self._write(data)

    elif type(payload) == types.GeneratorType:
      with self.lock:
        bytes = next(payload)
        first = True
        for chunk in payload:
          # This lets close() interrupt even a long running send.
          if self.server_terminate_request:
            break

          self._write(
            message_sender(bytes).fragment(first=first, mask=self.stream.always_mask)
          )
          bytes = chunk
          first = False

        self._write(
          message_sender(bytes).fragment(
            first=first, last=True, mask=self.stream.always_mask
          )
        )

    else:
      raise ValueError(f"Unsupported type '{type(payload)}' passed to send()")

  def _get_from_pending(self):
    """
    The SSL socket object provides the same interface
    as the socket interface but behaves differently.

    When data is sent over a SSL connection
    more data may be read than was requested from by
    the ws4py websocket object.

    In that case, the data may have been indeed read
    from the underlying real socket, but not read by the
    application which will expect another trigger from the
    manager's polling mechanism as if more data was still on the
    wire. This will happen only when new data is
    sent by the other peer which means there will be
    some delay before the initial read data is handled
    by the application.

    Due to this, we have to rely on a non-public method
    to query the internal SSL socket buffer if it has indeed
    more data pending in its buffer.

    Now, some people in the Python community
    `discourage <https://bugs.python.org/issue21430>`_
    this usage of the ``pending()`` method because it's not
    the right way of dealing with such use case. They advise
    `this approach <https://docs.python.org/dev/library/ssl.html#notes-on-non-blocking-sockets>`_
    instead. Unfortunately, this applies only if the
    application can directly control the poller which is not
    the case with the WebSocket abstraction here.

    We therefore rely on this `technic <http://stackoverflow.com/questions/3187565/select-and-ssl-in-python>`_
    which seems to be valid anyway.

    This is a bit of a shame because we have to process
    more data than what wanted initially.
    """
    data = b""
    pending = self.sock.pending()
    while pending:
      data += self.sock.recv(pending)
      pending = self.sock.pending()
    return data

  def once(self):
    """
    Performs the operation of reading from the underlying
    connection in order to feed the stream of bytes.

    Because this needs to support SSL sockets, we must always
    read as much as might be in the socket at any given time,
    however process expects to have itself called with only a certain
    number of bytes at a time. That number is found in
    self.reading_buffer_size, so we read everything into our own buffer,
    and then from there feed self.process.

    Then the stream indicates
    whatever size must be read from the connection since
    it knows the frame payload length.

    It returns `False` if an error occurred at the
    socket level or during the bytes processing. Otherwise,
    it returns `True`.
    """
    if self.terminated:
      logger.debug("WebSocket is already terminated")
      return False
    try:
      b = b""
      if self._is_secure:
        b = self._get_from_pending()
      if not b and not self.buf:
        b = self.sock.recv(self.reading_buffer_size)
      if not b and not self.buf:
        return False
      self.buf += b
    except (socket.error, OSError, pyOpenSSLError) as e:
      if hasattr(e, "errno") and e.errno == errno.EINTR:
        pass
      else:
        self.unhandled_error(e)
        return False
    else:
      # process as much as we can
      # the process will stop either if there is no buffer left
      # or if the stream is closed
      # only pass the requested number of bytes, leave the rest in the buffer
      requested = self.reading_buffer_size
      if not self.process(self.buf[:requested]):
        return False
      self.buf = self.buf[requested:]

    return True

  def terminate(self):
    """
    Completes the websocket by calling the `closed`
    method either using the received closing code
    and reason, or when none was received, using
    the special `1006` code.

    Finally close the underlying connection for
    good and cleanup resources by unsetting
    the `environ` and `stream` attributes.
    """
    s = self.stream

    if s is None:
      return

    try:
      if s.closing is None:
        self.closed(1006, "Going away")
      else:
        self.closed(s.closing.code, s.closing.reason)
    finally:
      self.client_terminated = self.server_terminated = True
      self.close_connection()

      # Cleaning up resources
      s._cleanup()
      self.stream = None
      self.environ = None

  def process(self, bytes):
    """Takes some bytes and process them through the
    internal stream's parser. If a message of any kind is
    found, performs one of these actions:

    * A closing message will initiate the closing handshake
    * Errors will initiate a closing handshake
    * A message will be passed to the ``received_message`` method
    * Pings will see pongs be sent automatically
    * Pongs will be passed to the ``ponged`` method

    The process should be terminated when this method
    returns ``False``.
    """
    s = self.stream

    if not bytes and self.reading_buffer_size > 0:
      return False

    with self.lock:
      self.reading_buffer_size = s.parser.send(bytes) or DEFAULT_READING_SIZE

    if s.closing is not None:
      logger.info(
        "Closing message received (%d) '%s'" % (s.closing.code, s.closing.reason)
      )
      if not self.server_terminated:
        self.close(s.closing.code, s.closing.reason)
      else:
        self.client_terminated = True
      return False

    if s.errors:
      for error in s.errors:
        logger.warn("Error message received (%d) '%s'" % (error.code, error.reason))
        self.close(error.code, error.reason)
      s.errors = []
      return False

    if s.has_message:
      self.received_message(s.message)
      if s.message is not None:
        s.message.data = None
        s.message = None
      return True

    if s.pings:
      for ping in s.pings:
        self._write(s.pong(ping.data))
      s.pings = []

    if s.pongs:
      for pong in s.pongs:
        self.ponged(pong)
      s.pongs = []

    return True

  def run(self):
    """
    Performs the operation of reading from the underlying
    connection in order to feed the stream of bytes.

    We start with a small size of two bytes to be read
    from the connection so that we can quickly parse an
    incoming frame header. Then the stream indicates
    whatever size must be read from the connection since
    it knows the frame payload length.

    Note that we perform some automatic opererations:

    * On a closing message, we respond with a closing
      message and finally close the connection
    * We respond to pings with pong messages.
    * Whenever an error is raised by the stream parsing,
      we initiate the closing of the connection with the
      appropiate error code.

    This method is blocking and should likely be run
    in a thread.
    """
    self.sock.setblocking(True)
    with Heartbeat(self, frequency=self.heartbeat_freq):
      s = self.stream

      try:
        self.opened()
        while not self.terminated:
          if not self.once():
            break
      except:
        logger.exception("Websocket connection was closed with an exception")
      finally:
        self.terminate()


class EchoWebSocket(WebSocket):
  def received_message(self, message):
    """
    Automatically sends back the provided ``message`` to
    its originating endpoint.
    """
    self.send(message.data, message.is_binary)
